import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.security.*;
import java.security.MessageDigest;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.Scanner;
import javax.crypto.spec.DESedeKeySpec;

public class Main {

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        try {
            while(true) {
                System.out.println("请选择操作：1-加密文件，2-解密文件，3-计算消息摘要，4-生成数字签名，5-验证数字签名，6-Diffie-Hellman密钥交换,7.退出系统");
                int choice = scanner.nextInt();
                scanner.nextLine(); // 消耗换行符

                switch (choice) {
                    case 1:
                        encryptFile(scanner);
                        break;
                    case 2:
                        decryptFile(scanner);
                        break;
                    case 3:
                        calculateDigest(scanner);
                        break;
                    case 4:
                        generateSignature(scanner);
                        break;
                    case 5:
                        verifySignature(scanner);
                        break;
                    case 6:
                        diffieHellmanKeyExchange();
                        break;
                    case 7:
                        System.out.println("退出系统！");

                        return;
                    default:
                        System.out.println("无效的选择");
                        break;
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            scanner.close();
        }
    }

    private static void encryptFile(Scanner scanner) throws Exception {
        System.out.println("输入要加密的文件路径：");
        String inputFilePath = scanner.nextLine();
        System.out.println("输入加密文件的保存路径：");
        String outputFilePath = scanner.nextLine();
        System.out.println("输入密钥保存路径：");
        String keyPath = scanner.nextLine();

        SecretKey key = generateKey("AES");
        encryptFile(new File(inputFilePath), new File(outputFilePath), key);
        saveKeyToFile(key, keyPath);
        System.out.println("密钥已保存到：" + keyPath);
    }

    private static void decryptFile(Scanner scanner) throws Exception {
        System.out.println("输入要解密的文件路径：");
        String inputFilePath = scanner.nextLine();
        System.out.println("输入解密文件的保存路径：");
        String outputFilePath = scanner.nextLine();
        System.out.println("输入密钥文件路径：");
        String keyPath = scanner.nextLine();

        SecretKey key = loadKeyFromFile(keyPath);
        decryptFile(new File(inputFilePath), new File(outputFilePath), key);
    }

    private static void calculateDigest(Scanner scanner) throws Exception {
        System.out.println("输入要计算摘要的信息：");
        String info = scanner.nextLine();

        byte[] digest = calculateDigest(info);
        System.out.println("信息摘要是：" + Base64.getEncoder().encodeToString(digest));
    }

    private static void generateSignature(Scanner scanner) throws Exception {
        System.out.println("输入要签名的信息：");
        String info = scanner.nextLine();

        System.out.println("输入私钥保存路径：");
        String privateKeyPath = scanner.nextLine();

        System.out.println("输入公钥保存路径：");
        String publicKeyPath = scanner.nextLine();

        KeyPair keyPair = generateKeyPair("DSA");
        byte[] signature = performSigning(info, "DSA", keyPair);
        System.out.println("签名是：" + Base64.getEncoder().encodeToString(signature));

        // 保存私钥和公钥
        savePrivateKeyToFile(keyPair.getPrivate(), privateKeyPath);
        savePublicKeyToFile(keyPair.getPublic(), publicKeyPath);
        System.out.println("密钥已保存。私钥路径：" + privateKeyPath + "，公钥路径：" + publicKeyPath);
    }

    private static void verifySignature(Scanner scanner) throws Exception {
        System.out.println("输入要验证的信息：");
        String info = scanner.nextLine();

        System.out.println("输入签名（Base64编码）：");
        String signatureBase64 = scanner.nextLine();

        System.out.println("输入公钥文件路径：");
        String publicKeyPath = scanner.nextLine();

        PublicKey publicKey = loadPublicKeyFromFile(publicKeyPath);
        boolean isValid = performVerification(info, "DSA", Base64.getDecoder().decode(signatureBase64), publicKey);
        System.out.println("签名验证结果：" + (isValid ? "有效" : "无效"));
    }

    private static void diffieHellmanKeyExchange() throws Exception {
        KeyPairGenerator aliceKpairGen = KeyPairGenerator.getInstance("DH");
        aliceKpairGen.initialize(512);
        KeyPair aliceKpair = aliceKpairGen.generateKeyPair();

        KeyAgreement aliceKeyAgree = KeyAgreement.getInstance("DH");
        aliceKeyAgree.init(aliceKpair.getPrivate());
        DHPublicKey alicePublicKey = (DHPublicKey)aliceKpair.getPublic();
        System.out.println("Alice的公钥：" + Base64.getEncoder().encodeToString(aliceKpair.getPublic().getEncoded()));

        DHParameterSpec dhParamSpec = ((DHPublicKey) aliceKpair.getPublic()).getParams();
        KeyPairGenerator bobKpairGen = KeyPairGenerator.getInstance("DH");
        bobKpairGen.initialize(dhParamSpec);
        KeyPair bobKpair = bobKpairGen.generateKeyPair();
        DHPublicKey bobPublicKey = (DHPublicKey)bobKpair.getPublic();
        KeyAgreement bobKeyAgree = KeyAgreement.getInstance("DH");
        bobKeyAgree.init(bobKpair.getPrivate());
        System.out.println("Bob的公钥：" + Base64.getEncoder().encodeToString(bobKpair.getPublic().getEncoded()));

        aliceKeyAgree.doPhase(bobPublicKey, true);
        bobKeyAgree.doPhase(alicePublicKey, true);
        SecretKey aliceDesKey = bobKeyAgree.generateSecret("DES");
        SecretKey bobDesKey = aliceKeyAgree.generateSecret("DES");

        System.out.println("Alice和Bob共享的DES密钥：" + Base64.getEncoder().encodeToString(aliceDesKey.getEncoded()));
        System.out.println("Bob和Alice共享的DES密钥：" + Base64.getEncoder().encodeToString(bobDesKey.getEncoded()));
    }

    public static SecretKey generateKey(String algorithm) throws Exception {
        KeyGenerator keyGenerator = KeyGenerator.getInstance(algorithm);
        keyGenerator.init(256);
        return keyGenerator.generateKey();
    }

    public static KeyPair generateKeyPair(String alg) throws Exception {
        KeyPairGenerator kg = KeyPairGenerator.getInstance(alg);
        kg.initialize(512);
        return kg.genKeyPair();
    }

    public static byte[] calculateDigest(String message) throws Exception {
        MessageDigest digest = MessageDigest.getInstance("SHA-256");
        return digest.digest(message.getBytes());
    }

    public static byte[] performSigning(String s, String alg, KeyPair keyPair) throws Exception {
        Signature sign = Signature.getInstance(alg);
        PrivateKey privateKey = keyPair.getPrivate();
        sign.initSign(privateKey);
        sign.update(s.getBytes());
        return sign.sign();
    }

    public static boolean performVerification(String s, String alg, byte[] signature, PublicKey publicKey) throws Exception {
        Signature sign = Signature.getInstance(alg);
        sign.initVerify(publicKey);
        sign.update(s.getBytes());
        return sign.verify(signature);
    }

    public static void encryptFile(File inputFile, File outputFile, SecretKey key) throws Exception {
        Cipher cipher = Cipher.getInstance("AES");
        cipher.init(Cipher.ENCRYPT_MODE, key);

        try (FileInputStream inputStream = new FileInputStream(inputFile);
             FileOutputStream outputStream = new FileOutputStream(outputFile)) {
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                byte[] output = cipher.update(buffer, 0, bytesRead);
                if (output != null) {
                    outputStream.write(output);
                }
            }
            byte[] output = cipher.doFinal();
            if (output != null) {
                outputStream.write(output);
            }
        }
    }

    public static void decryptFile(File inputFile, File outputFile, SecretKey key) throws Exception {
        Cipher cipher = Cipher.getInstance("AES");
        cipher.init(Cipher.DECRYPT_MODE, key);

        try (FileInputStream inputStream = new FileInputStream(inputFile);
             FileOutputStream outputStream = new FileOutputStream(outputFile)) {
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                byte[] output = cipher.update(buffer, 0, bytesRead);
                if (output != null) {
                    outputStream.write(output);
                }
            }
            byte[] output = cipher.doFinal();
            if (output != null) {
                outputStream.write(output);
            }
        }
    }

    public static void saveKeyToFile(SecretKey key, String filePath) throws IOException {
        try (FileOutputStream fos = new FileOutputStream(filePath)) {
            fos.write(Base64.getEncoder().encode(key.getEncoded()));
        }
    }

    public static SecretKey loadKeyFromFile(String filePath) throws Exception {
        try (FileInputStream fis = new FileInputStream(filePath)) {
            byte[] keyBytes = new byte[fis.available()];
            fis.read(keyBytes);
            return new SecretKeySpec(Base64.getDecoder().decode(keyBytes), "AES");
        }
    }
    public static void savePrivateKeyToFile(PrivateKey key, String filePath) throws IOException {
        try (FileOutputStream fos = new FileOutputStream(filePath)) {
            fos.write(key.getEncoded());
        }
    }

    public static void savePublicKeyToFile(PublicKey key, String filePath) throws IOException {
        try (FileOutputStream fos = new FileOutputStream(filePath)) {
            fos.write(key.getEncoded());
        }
    }

    public static PublicKey loadPublicKeyFromFile(String filePath) throws Exception {
        try (FileInputStream fis = new FileInputStream(filePath)) {
            byte[] keyBytes = new byte[fis.available()];
            fis.read(keyBytes);
            KeyFactory keyFactory = KeyFactory.getInstance("DSA");
            return keyFactory.generatePublic(new X509EncodedKeySpec(keyBytes));
        }
    }
}

